package com.sunchp.nio.sample.nio2;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.util.Iterator;

public class SelectThread extends Thread {
    private final Logger LOGGER = LoggerFactory.getLogger(SelectThread.class);

    protected volatile boolean stopped = false;
    private Selector selector;
    private ServerSocketChannel serverSocketChannel;

    public SelectThread() throws IOException {
        this.selector = Selector.open();
        this.serverSocketChannel = ServerSocketChannel.open();
        this.serverSocketChannel.configureBlocking(false);
        this.serverSocketChannel.socket().bind(new InetSocketAddress(8080));
        this.serverSocketChannel.register(selector, SelectionKey.OP_ACCEPT);
    }

    public void run() {
        try {
            while (!stopped) {
                select();
            }

            for (SelectionKey selectionKey : selector.keys()) {
                cleanupSelectionKey(selectionKey);
            }
        } catch (Throwable t) {
            LOGGER.error("run() exiting due to uncaught error", t);
        } finally {
            try {
                selector.close();
            } catch (IOException e) {
                LOGGER.error("Got an IOException while closing selector!", e);
            }

            stopped = true;
        }
    }

    private void select() {
        try {
            selector.select();

            Iterator<SelectionKey> selectedKeys = selector.selectedKeys().iterator();
            while (selectedKeys.hasNext()) {
                SelectionKey key = selectedKeys.next();
                selectedKeys.remove();

                if (!key.isValid()) {
                    cleanupSelectionKey(key);
                    continue;
                }

                if (key.isAcceptable()) {
                    handleAccept();
                } else if (key.isReadable()) {
                    handleRead(key);
                } else if (key.isWritable()) {
                    handleWrite(key);
                } else {
                    LOGGER.warn("Unexpected state in select! " + key.interestOps());
                }
            }
        } catch (IOException e) {
            LOGGER.warn("Got an IOException while selecting!", e);
        }
    }

    private void cleanupSelectionKey(SelectionKey key) {
        FrameBuffer buffer = (FrameBuffer) key.attachment();
        if (buffer != null) {
            buffer.close();
        }
        key.cancel();
    }

    private FrameBuffer createFrameBuffer(SocketChannel socketChannel, SelectionKey selectionKey) {
        return new FrameBuffer(socketChannel, selectionKey);
    }

    private void handleAccept() throws IOException {
        SelectionKey clientKey = null;
        SocketChannel client = null;
        try {
            client = serverSocketChannel.accept();
            client.configureBlocking(false);
            clientKey = client.register(selector, SelectionKey.OP_READ);

            FrameBuffer frameBuffer = createFrameBuffer(client, clientKey);
            clientKey.attach(frameBuffer);
        } catch (TransportException tte) {
            LOGGER.error("Exception trying to accept!", tte);
            if (clientKey != null) {
                cleanupSelectionKey(clientKey);
            }

            if (client != null) {
                client.close();
            }
        }
    }

    private void handleRead(SelectionKey key) {
        FrameBuffer buffer = (FrameBuffer) key.attachment();
        if (!buffer.read()) {
            cleanupSelectionKey(key);
            return;
        }

        if (buffer.isFrameFullyRead()) {
            buffer.invoke();
        }
    }

    private void handleWrite(SelectionKey key) {
        FrameBuffer buffer = (FrameBuffer) key.attachment();
        if (!buffer.write()) {
            cleanupSelectionKey(key);
        }
    }
}
